#include "nets.hpp"

CustomDtaset::CustomDtaset(at::Tensor input_data, at::Tensor label_data)
{
    inputs_ = input_data;
    labels_ = label_data;
}

torch::data::Example<> CustomDtaset::get(size_t index)
{
    torch::Tensor Input = inputs_[index];
    torch::Tensor Label = labels_[index];
    return {Input.clone(), Label.clone()};
}

torch::optional<size_t> CustomDtaset::size() const { return labels_.size(0); }

Net::Net(int in_dim, int out_dim)
{
    k = register_parameter("k", torch::zeros({in_dim, out_dim}), true);
    b = register_parameter("b", torch::zeros(out_dim), true);
}

at::Tensor Net::forward(at::Tensor x) { return torch::addmm(b, x, k); }
